// Copyright(C) 2023 InfiniFlow, Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

module infinity_core:thrift_server.impl;

import :thrift_server;
import :infinity_thrift_service;
import :infinity_thrift_types;
import :infinity_exception;

import std;
import third_party;

namespace infinity {

class CustomTTransport final : public apache::thrift::protocol::TTransport {
public:
    CustomTTransport() : apache::thrift::protocol::TTransport(nullptr) {}
};

class CustomerBufferedTransportFactory : public apache::thrift::transport::TTransportFactory {
public:
    CustomerBufferedTransportFactory() = default;

    ~CustomerBufferedTransportFactory() override = default;

    std::shared_ptr<apache::thrift::protocol::TTransport> getTransport(std::shared_ptr<apache::thrift::protocol::TTransport> trans) override {
        auto configuration = std::make_shared<apache::thrift::TConfiguration>();
        configuration->setRecursionLimit(256);
        trans->setConfiguration(configuration);
        return std::shared_ptr<apache::thrift::protocol::TTransport>(new apache::thrift::transport::TBufferedTransport(trans));
    }
};

class InfinityServiceCloneFactory final : public infinity_thrift_rpc::InfinityServiceIfFactory {
public:
    ~InfinityServiceCloneFactory() final = default;

    infinity_thrift_rpc::InfinityServiceIf *getHandler(const apache::thrift::TConnectionInfo &connInfo) final { return new InfinityThriftService; }

    void releaseHandler(infinity_thrift_rpc::InfinityServiceIf *handler) final { delete handler; }
};

// Thrift server
#define THRIFT_SERVER_TYPE 0

#if THRIFT_SERVER_TYPE == 2

void ThreadedThriftServer::Init(const std::string &server_address, i32 port_no) {

    fmt::print("API server listen on {}: {}\n", server_address, port_no);
    //    std::cout << "API server listen on: " << server_address << ": " << port_no << std::endl;
    std::shared_ptr<TBinaryProtocolFactory> binary_protocol_factory = std::make_shared<TBinaryProtocolFactory>();
    binary_protocol_factory->setStrict(true, true);
    server = std::make_unique<TThreadedServer>(
        std::make_shared<infinity_thrift_rpc::InfinityServiceProcessorFactory>(std::make_shared<InfinityServiceCloneFactory>()),
        std::make_shared<TServerSocket>(server_address, port_no),
        std::make_shared<TBufferedTransportFactory>(),
        binary_protocol_factory);
}

void ThreadedThriftServer::Start() {
    if (started_) {
        return;
    }
    started_ = true;
    server->serve();
}

void ThreadedThriftServer::Shutdown() {
    if (started_) {
        server->stop();
        started_ = false;
    }
}

#elif THRIFT_SERVER_TYPE == 0

void PoolThriftServer::Init(const std::string &server_address, i32 port_no, i32 pool_size) {

    auto server_socket = std::make_shared<apache::thrift::transport::TServerSocket>(server_address, port_no);

    auto protocol_factory = std::make_shared<apache::thrift::protocol::TBinaryProtocolFactory>();
    //    std::shared_ptr<TCompactProtocolFactory> protocol_factory = std::make_shared<TCompactProtocolFactory>();

    auto threadFactory = std::make_shared<apache::thrift::concurrency::ThreadFactory>();

    std::shared_ptr<apache::thrift::concurrency::ThreadManager> threadManager =
        apache::thrift::concurrency::ThreadManager::newSimpleThreadManager(pool_size);
    threadManager->threadFactory(threadFactory);
    threadManager->start();

    fmt::print("API server(for Infinity-SDK) listen on {}: {}, connection limit: {}\n", server_address, port_no, pool_size);
    //    std::cout << "API server listen on: " << server_address << ": " << port_no << ", thread pool: " << pool_size << std::endl;

    server = std::make_unique<apache::thrift::server::TThreadPoolServer>(
        std::make_shared<infinity_thrift_rpc::InfinityServiceProcessorFactory>(std::make_shared<InfinityServiceCloneFactory>()),
        server_socket,
        std::make_shared<CustomerBufferedTransportFactory>(),
        protocol_factory,
        threadManager);

    initialized_ = true;
}

std::thread PoolThriftServer::Start() {
    if (!initialized_) {
        UnrecoverableError("Thrift server is not initialized");
    }
    {
        if (auto expect = ThriftServerStatus::kStopped; !status_.compare_exchange_strong(expect, ThriftServerStatus::kRunning)) {
            UnrecoverableError(fmt::format("Thrift server in unexpected state: {}", u8(expect)));
        }
    }
    return std::thread([this] {
        server->serve();

        status_.store(ThriftServerStatus::kStopped);
        status_.notify_one();
    });
}

void PoolThriftServer::Shutdown() {
    {
        auto expected = ThriftServerStatus::kRunning;
        if (!status_.compare_exchange_strong(expected, ThriftServerStatus::kStopping)) {
            if (status_ == ThriftServerStatus::kStopped) {
                return;
            } else {
                UnrecoverableError(fmt::format("Thrift server in unexpected state: {}", u8(expected)));
            }
        }
    }
    server->stop();

    status_.wait(ThriftServerStatus::kStopping);
}

#elif THRIFT_SERVER_TYPE == 1

void NonBlockPoolThriftServer::Init(const std::string &server_address, i32 port_no, i32 pool_size) {

    std::shared_ptr<ThreadFactory> thread_factory = std::make_shared<ThreadFactory>();
    service_handler_ = std::make_shared<InfinityThriftService>();
    std::shared_ptr<infinity_thrift_rpc::InfinityServiceProcessor> service_processor =
        std::make_shared<infinity_thrift_rpc::InfinityServiceProcessor>(service_handler_);
    std::shared_ptr<TProtocolFactory> protocol_factory = std::make_shared<TBinaryProtocolFactory>();

    std::shared_ptr<ThreadManager> threadManager = ThreadManager::newSimpleThreadManager(pool_size);
    threadManager->threadFactory(thread_factory);
    threadManager->start();

    fmt::print("Non-block API server listen on {}: {}, connection limit: {}\n", server_address, port_no, pool_size);
    //    std::cout << "Non-block API server listen on: " << server_address << ": " << port_no << ", thread pool: " << pool_size << std::endl;

    std::shared_ptr<TNonblockingServerSocket> non_block_socket = std::make_shared<TNonblockingServerSocket>(server_address, port_no);

    //    server_thread_ = thread_factory->newThread(std::shared_ptr<TServer>(
    //        new TNonblockingServer(serviceProcessor, protocolFactory, nbSocket1, threadManager)));

    server_thread_ =
        thread_factory->newThread(std::make_shared<TNonblockingServer>(service_processor, protocol_factory, non_block_socket, threadManager));

    //    server =
    //    std::make_unique<TThreadPoolServer>(std::make_shared<InfinityServiceProcessorFactory>(std::make_shared<InfinityServiceCloneFactory>()),
    //                                           std::make_shared<TServerSocket>(port_no),
    //                                           std::make_shared<TBufferedTransportFactory>(),
    //                                           std::make_shared<TBinaryProtocolFactory>(),
    //                                           threadManager);
}

void NonBlockPoolThriftServer::Start() {
    if (started_) {
        return;
    }
    started_ = true;
    server_thread_->start();
}

void NonBlockPoolThriftServer::Shutdown() {
    if (started_) {
        server_thread_->join();
        started_ = false;
    }
}

#endif

} // namespace infinity
